package com.smmisia.auth.config;

import com.smmisia.auth.aop.Encrypted;
import com.smmisia.auth.utls.AESUtil;
import org.springframework.core.MethodParameter;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.RequestBodyAdviceAdapter;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;

/**
 * @author smmisia
 * @date 2024/4/16 16:44
 * 解密请求体
 */
@RestControllerAdvice
public class SecretRequestAdvice extends RequestBodyAdviceAdapter {

    private static final String SECRET_KEY = "smmisia"; // 替换为你的密钥

    @Override
    public boolean supports(MethodParameter methodParameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        // 判断是否需要解密（通过注解判断）
        return methodParameter.hasMethodAnnotation(Encrypted.class);
    }

    @Override
    public HttpInputMessage beforeBodyRead(HttpInputMessage inputMessage, MethodParameter parameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) throws IOException {
        try {
            // 解密请求体
            String decryptedBody = decrypt(inputMessage.getBody());
            return new DecryptedHttpInputMessage(inputMessage, decryptedBody);
        } catch (Exception e) {
            throw new RuntimeException("Failed to decrypt request body", e);
        }
    }

    private String decrypt(InputStream inputStream) throws Exception {
        // 读取请求体并解密
        String encryptedBody = new String(inputStream.readAllBytes(), StandardCharsets.UTF_8);
        return AESUtil.decrypt(encryptedBody, SECRET_KEY);
    }

    private static class DecryptedHttpInputMessage implements HttpInputMessage {
        private final HttpInputMessage delegate;
        private final String decryptedBody;

        public DecryptedHttpInputMessage(HttpInputMessage delegate, String decryptedBody) {
            this.delegate = delegate;
            this.decryptedBody = decryptedBody;
        }

        @Override
        public InputStream getBody() {
            return new ByteArrayInputStream(decryptedBody.getBytes(StandardCharsets.UTF_8));
        }

        @Override
        public HttpHeaders getHeaders() {
            return delegate.getHeaders();
        }
    }
}
